package com.zhny.test;

import org.apache.parquet.Strings;
import org.apache.spark.mllib.linalg.Vector;
import org.apache.spark.mllib.linalg.Vectors;
import scala.Tuple2;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.mllib.classification.NaiveBayes;
import org.apache.spark.mllib.classification.NaiveBayesModel;
import org.apache.spark.mllib.regression.LabeledPoint;
import org.apache.spark.mllib.util.MLUtils;
// $example off$
import org.apache.spark.SparkConf;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
//贝叶斯
public class NaiveBayesUtil {
    public static void exc(String dataFilePath, String resultFilePath) {
        SparkConf sparkConf = new SparkConf().setAppName("JavaNaiveBayesExample").setMaster("local[1]");
        sparkConf.set("spark.driver.allowMultipleContexts", "true");
        JavaSparkContext jsc = new JavaSparkContext(sparkConf);
        // $example on$

        String path = "";
        if (Strings.isNullOrEmpty(dataFilePath)) {
            path = "src/main/resources/data/NaiveBayesData.txt";
        } else {
            path = dataFilePath;
        }

        JavaRDD<LabeledPoint> inputData = MLUtils.loadLibSVMFile(jsc.sc(), path).toJavaRDD();
        JavaRDD<LabeledPoint>[] tmp = inputData.randomSplit(new double[]{0.6, 0.4});
        JavaRDD<LabeledPoint> training = tmp[0]; // training set
        JavaRDD<LabeledPoint> test = tmp[1]; // test set

        NaiveBayesModel model = NaiveBayes.train(training.rdd(), 1.0);

        JavaPairRDD<Double, Double> predictionAndLabel =
                test.mapToPair(p -> new Tuple2<>(model.predict(p.features()), p.label()));
        double accuracy =
                predictionAndLabel.filter(pl -> pl._1().equals(pl._2())).count() / (double) test.count();

        // Save and load model
//        model.save(jsc.sc(), "src/main/resources/data/modal/NaiveBayesModel");
//        NaiveBayesModel sameModel = NaiveBayesModel.load(jsc.sc(), "src/main/resources/data/modal/NaiveBayesModel");
        // $example off$

        double[] d1 = new double[] { 78, 36, 66, 43, 45, 40 };
        Vector v1 = Vectors.dense(d1);

        FileWriter fos = null;
        try {
            fos = new FileWriter(new File(resultFilePath));
//            fos.write(String.valueOf(new Double(model.predict(v1)).longValue()) + "\n");
            fos.write(String.valueOf(new Double(accuracy)) + "\n");

            fos.flush();
            fos.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        jsc.stop();
    }
}
